#ifndef UHM_SOLVER_BASE_PETSC_HXX_
#define UHM_SOLVER_BASE_PETSC_HXX_

#include "petscksp.h"

namespace uhm {

  class PETSC_Solver_Base_  {
  protected:
    Int_ _d_nz, _o_nz, _alloc;
    Double_ _rtol, _atol; 
    std::string _pctype;

  public:
    PETSC_Solver_Base_() : _rtol(1.0e-7), 
                           _atol(1.0e-10), 
                           _pctype("ilu"),
                           _d_nz(0),
                           _o_nz(0),
                           _alloc(100) { }
    virtual~PETSC_Solver_Base_()  { }

    virtual void set_alloc_scale(Int_ alloc) { _alloc = alloc; }
    virtual void set_tolerence(Double_ rtol, Double_ atol) { _rtol = rtol; _atol = atol;  }
    virtual void set_pctype(Char_ *pctype) { _pctype = pctype; }
    // - Type
    // PCNONE, PCJACOBI, PCSOR, PCLU, PCSHELL, PCBJACOBI, PCMG, PCILU
    // PCASM, PCKSP, PCCOMPOSITE, PCSPI, PCCHOLESKY, PCSAMG, PCHYPRE and so on
    //
    // see petscpc.h
    
    virtual Int_ assemble() { return false; }    
  };

  template<class Id_Nd_, class Id_El_, class Uhm_>
  class PETSC_Solver_ : public Direct_Solver_<Id_Nd_,Id_El_,Uhm_>,
                        public PETSC_Solver_Base_ {
  private:
    Mat _A;
    Vec _x, _b, _r;
    KSP _ksp;

    Map_<Id_Nd_> _nz;
    std::vector< std::vector<Int_> > _idx;

  public:

    // typedef is not inherited, explicitly bring Solver typeder here...
    typedef Solver_<Id_Nd_,Id_El_,Uhm_> _Solver_;

    typedef Mesh_<Id_Nd_,Id_El_,Uhm_> _Mesh_;
    typedef Tree_<typename _Mesh_::Element_*,UHM_TREE_MAX_CHILDREN> _Tree_;
    typedef Cache_<Id_Nd_,Uhm_> _Cache_;

    PETSC_Solver_() 
      : Direct_Solver_<Id_Nd_,Id_El_,Uhm_>(),
        PETSC_Solver_Base_()   { } 

    PETSC_Solver_(Int_ type, Int_ uplo, Int_ sym,
                  Int_ rhs, Int_ blk)  
      : Direct_Solver_<Id_Nd_,Id_El_,Uhm_>(type, uplo, sym, rhs, blk),
        PETSC_Solver_Base_()   { } 
    
    PETSC_Solver_(Int_ type, Int_ rhs, Int_ blk) 
      : Direct_Solver_<Id_Nd_,Id_El_,Uhm_>(type, rhs, blk),
        PETSC_Solver_Base_()   { }
    
    virtual~PETSC_Solver_() { }

    //using Direct_Solver_<Id_Nd_,Id_El_,Uhm_>::get_flop;
    //using Direct_Solver_<Id_Nd_,Id_El_,Uhm_>::decompose;
    //using Direct_Solver_<Id_Nd_,Id_El_,Uhm_>::updecompose;
    //using Direct_Solver_<Id_Nd_,Id_El_,Uhm_>::check;
    //using Direct_Solver_<Id_Nd_,Id_El_,Uhm_>::refine;

    virtual Int_ create() {
      time_in("uhm:: petsc:: create",UHM_TIME_LEVEL_FRONT);

      // Cutting a tree and setting dofs and offsets
      Direct_Solver_<Id_Nd_,Id_El_,Uhm_>::create();
      
      // Sparcity structure
      time_in("uhm:: petsc:: create:: Identify Sparcity Patterns",UHM_TIME_LEVEL_INTERNAL);
      _nz.push_begin(this->mesh().get_n_nodes());
      for (Int_ i=0;i<this->tnds().size();++i) {
        Map_<Id_Nd_> &schur = this->tnds()[i]->get_data()->schur();
        for (Int_ j=0;j<schur.get_size();++j) 
          _nz.push(schur.id(j));
      }
      _nz.push_end();
      _nz.unique();

      // So far idx and dof in _nz are set up right, but dofs will be deleted 
      // then set it up as entire row or column info
      _nz.delete_local_dof();
      for (Int_ i=0;i<this->tnds().size();++i) {
        Map_<Id_Nd_> &schur = this->tnds()[i]->get_data()->schur();
        // Sum
        for (Int_ j=0;j<schur.get_size();++j) {
          Id_Nd_ nd; 
          _nz.update_local_dof(schur.id(j));
        }
      }

      // _d_nz is max of row or column size
      this->_d_nz = 0;
      for (Int_ i=0;i<_nz.get_size();++i) {
        this->_d_nz = (this->_d_nz < _nz.id(i).get_dof() ? _nz.id(i).get_dof() : this->_d_nz);
      }
      time_out(UHM_TIME_LEVEL_INTERNAL);

      PetscErrorCode ierr;

      time_in("uhm:: petsc:: create:: PETSC Matrix Creating",UHM_TIME_LEVEL_INTERNAL);
      ierr = MatCreate(PETSC_COMM_WORLD,&_A);CHKERRQ(ierr);
      ierr = MatSetSizes(_A, PETSC_DECIDE,PETSC_DECIDE,
                         _nz.dof(), _nz.dof());CHKERRQ(ierr);  
      // - Options
      // -mat_type seqaij   - AIJ type, uses MatCreateSeqAIJ()
      // -mat_type mpiaij   - AIJ type, uses MatCreateMPIAIJ()
      // -mat_type seqdense - dense type, uses MatCreateSeqDense()
      // -mat_type mpidense - dense type, uses MatCreateMPIDense()
      // -mat_type seqbaij  - block AIJ type, uses MatCreateSeqBAIJ()
      // -mat_type mpibaij  - block AIJ type, uses MatCreateMPIBAIJ()

      // As default, MATSEQAIJ or MATMPIAIJ
      //ierr = MatSetType(_A, MATSEQAIJ);CHKERRQ(ierr);
      //ierr = MatSetOption(_A, MAT_ROW_ORIENTED, PETSC_TRUE);
      //ierr = MatSetOption(_A, MAT_STRUCTURALLY_SYMMETRIC, PETSC_TRUE);CHKERRQ(ierr);
      ierr = MatSetFromOptions(_A);CHKERRQ(ierr);
      time_out(UHM_TIME_LEVEL_INTERNAL);

      time_in("uhm:: petsc:: create:: PETSC Vector Creating",UHM_TIME_LEVEL_INTERNAL);
      ierr = VecCreate(PETSC_COMM_WORLD, &_b);CHKERRQ(ierr);
      ierr = VecSetSizes(_b,PETSC_DECIDE, _nz.dof());CHKERRQ(ierr);
      ierr = VecSetFromOptions(_b);CHKERRQ(ierr);

      ierr = VecDuplicate(_b, &_x);CHKERRQ(ierr);
      ierr = VecDuplicate(_x, &_r);CHKERRQ(ierr);
      time_out(UHM_TIME_LEVEL_INTERNAL);
      time_out(UHM_TIME_LEVEL_FRONT);

      return true;
    }

    virtual Int_ clear() {
      Direct_Solver_<Id_Nd_,Id_El_,Uhm_>::clear();

      PetscErrorCode ierr;
      
      time_in("uhm:: petsc:: clear",UHM_TIME_LEVEL_FRONT);
      ierr = MatDestroy(&_A);CHKERRQ(ierr);
      ierr = VecDestroy(&_x);CHKERRQ(ierr);
      ierr = VecDestroy(&_b);CHKERRQ(ierr);
      ierr = VecDestroy(&_r);CHKERRQ(ierr);
      time_out(UHM_TIME_LEVEL_FRONT);

      return true;
    }

    virtual Int_ solve() {
      time_in("uhm:: petsc:: solve",UHM_TIME_LEVEL_FRONT);
      time_in("uhm:: petsc:: solve:: Solving Phase 0, 1",UHM_TIME_LEVEL_INTERNAL);
      { // ** Solve Phase 0, 1
#pragma omp parallel
        {
#pragma omp single nowait
          {
		  
            for (Int_ i=0;i<this->tnds().size();++i) {
#pragma omp task firstprivate(i)
              {
                UHM_ERROR(op_leaf_to_me_par<_Solver_>(this, this->tnds()[i],
                                                      &op_uhm_direct_solve_0_x<_Solver_>),
                          UHM_ERROR_DIRECT_SOLVE_ZERO);
                UHM_ERROR(op_leaf_to_me_par<_Solver_>(this, this->tnds()[i],
                                                      &op_uhm_direct_solve_1_x<_Solver_>),
                          UHM_ERROR_DIRECT_SOLVE_ONE);
              }
            }
          }
		  
        }
      }
      time_out(UHM_TIME_LEVEL_INTERNAL);

      time_in("uhm:: petsc:: solve:: PETSC Solving a Schur System",UHM_TIME_LEVEL_INTERNAL);
      if (this->tnds().size() > 1) {
        PetscErrorCode ierr;
        
        // - Options
        // SAME_PRECONDITIONER - Pmat is identical during successive linear solves
        // SAME_NONZERO_PATTERN - Pmat has same structure
        // DIFFERENT_NONZERO_PATTERN - different structure
        
        // Default is GMRES
        time_in("uhm:: petsc:: solve:: PETSC Creating Krylov Solver",UHM_TIME_LEVEL_INTERNAL);
        ierr = KSPCreate(PETSC_COMM_WORLD, &_ksp);CHKERRQ(ierr);
        ierr = KSPSetOperators(_ksp, _A, _A, SAME_PRECONDITIONER);CHKERRQ(ierr);
        ierr = KSPSetTolerances(_ksp, this->_rtol, this->_atol,
                                PETSC_DEFAULT, PETSC_DEFAULT);CHKERRQ(ierr);
        ierr = KSPSetInitialGuessNonzero(_ksp, PETSC_TRUE);CHKERRQ(ierr);

        PC pc;
        ierr = KSPGetPC(_ksp, &pc);CHKERRQ(ierr);
        ierr = PCSetType(pc, this->_pctype.c_str());CHKERRQ(ierr);

        ierr = KSPSetFromOptions(_ksp);CHKERRQ(ierr);
        
        // Setting up Solver
        ierr = KSPSetUp(_ksp);
        time_out(UHM_TIME_LEVEL_INTERNAL);
          
        for (Int_ n=0;n<this->get_rhs();++n) {
        
          time_in("uhm:: petsc:: solve:: PETSC Assembling Right Hand Side (one-by-one)",
                  UHM_TIME_LEVEL_INTERNAL);
          for (Int_ i=0;i<this->tnds().size();++i) {
            // get_flat_buffer needs shift
            Uhm_ &uhm = this->tnds()[i]->get_data()->uhm();
            ierr = VecSetValues(_b, 
                                uhm.XB().get_m(), &_idx[i][0], 
                                (const PetscScalar*)(uhm.XB().get_flat_buffer(n)),
                                ADD_VALUES);CHKERRQ(ierr);
          }

          ierr = VecAssemblyBegin(_b);CHKERRQ(ierr);
          ierr = VecAssemblyEnd  (_b);CHKERRQ(ierr);
          time_out(UHM_TIME_LEVEL_INTERNAL);

          time_in("uhm:: petsc:: solve:: PETSC Solving (Krylov)",UHM_TIME_LEVEL_INTERNAL);
          ierr = KSPSolve(_ksp, _b, _x);CHKERRQ(ierr);
          ierr = KSPView(_ksp, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
          PetscInt it;

          PetscInt its;
          KSPConvergedReason reason;
          ierr = KSPGetConvergedReason(_ksp, &reason);CHKERRQ(ierr);

          PetscReal rnorm;
          ierr = KSPGetResidualNorm(_ksp, &rnorm);CHKERRQ(ierr);
          ierr = PetscPrintf(PETSC_COMM_WORLD,"\nPETSC Residual = %E\n", rnorm);
          if (reason < 0) {
            ierr = PetscPrintf(PETSC_COMM_WORLD,"Diverged reason = %d\n\n", reason);
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_NULL                = -2\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_ITS                 = -3\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_DTOL                = -4\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_BREAKDOWN           = -5\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_BREAKDOWN_BICG      = -6\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_NONSYMMETRIC        = -7\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_INDEFINITE_PC       = -8\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_NAN                 = -9\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_DIVERGED_INDEFINITE_MAT      = -10\n");
          } else {
            ierr = KSPGetIterationNumber(_ksp, &its);
            ierr = PetscPrintf(PETSC_COMM_WORLD,"Converged reason = %d, # Iteration = %d\n\n", 
                               reason, its);
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_RTOL_NORMAL        =  1\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_RTOL               =  2\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_ATOL               =  3\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_ITS                =  4\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_CG_NEG_CURVE       =  5\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_CG_CONSTRAINED     =  6\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_STEP_LENGTH        =  7\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_HAPPY_BREAKDOWN    =  8\n");
            ierr = PetscPrintf(PETSC_COMM_WORLD,"KSP_CONVERGED_ATOL_NORMAL        =  9\n");
          }
          time_out(UHM_TIME_LEVEL_INTERNAL);

          // Solution out
          time_in("uhm:: petsc:: solve:: PETSC Branching Solutions",UHM_TIME_LEVEL_INTERNAL);
#pragma omp parallel for schedule(static)
          for (Int_ i=0;i<this->tnds().size();++i) {
            Uhm_ &uhm = this->tnds()[i]->get_data()->uhm();
            ierr = VecGetValues(_x,
                                uhm.XB().get_m(), &_idx[i][0],
                                (PetscScalar*)(uhm.XB().get_flat_buffer(n)));CHKERRQ(ierr);
          }
          time_out(UHM_TIME_LEVEL_INTERNAL);
        }
        time_in("uhm:: petsc:: solve:: PETSC Destroying Solver (Krylov)",UHM_TIME_LEVEL_INTERNAL);
        ierr = KSPDestroy(&_ksp);CHKERRQ(ierr);
        time_out(UHM_TIME_LEVEL_INTERNAL);
      }
      time_out(UHM_TIME_LEVEL_INTERNAL);

      time_in("uhm:: petsc:: solve:: Solving Phase 2",UHM_TIME_LEVEL_INTERNAL);
      {
#pragma omp parallel
        {
#pragma omp single nowait
          {

            for (Int_ i=0;i<this->tnds().size();++i) {
#pragma omp task firstprivate(i)
              {
                UHM_ERROR(op_me_to_leaf_par<_Solver_>(this, this->tnds()[i],
                                                      &op_uhm_direct_solve_2_x<_Solver_>),
                          UHM_ERROR_DIRECT_SOLVE_TWO);
              }
            }
          }
        }
      }
      time_out(UHM_TIME_LEVEL_INTERNAL);
      time_out(UHM_TIME_LEVEL_FRONT);
      return true;
    }

    virtual Int_ assemble() {
      time_in("uhm:: petsc:: assemble",UHM_TIME_LEVEL_FRONT);
      PetscErrorCode ierr;

      // Right now this is sequential... 
      time_in("uhm:: petsc:: assemble:: Preallocation",UHM_TIME_LEVEL_INTERNAL);
      ierr = MatMPIAIJSetPreallocation(_A, 
                                       this->_d_nz*this->_alloc, PETSC_NULL, 
                                       this->_o_nz*this->_alloc, PETSC_NULL);CHKERRQ(ierr);
      ierr = MatSeqAIJSetPreallocation(_A, 
                                       this->_d_nz*this->_alloc, PETSC_NULL);CHKERRQ(ierr);
      time_out(UHM_TIME_LEVEL_INTERNAL);

      // Creating bijection map local2global
      time_in("uhm:: petsc:: assemble:: Assembling Matrices",UHM_TIME_LEVEL_INTERNAL);
      _idx.clear();
      _idx.reserve(this->tnds().size());
      for (Int_ i=0;i<this->tnds().size();++i) {
        _idx.push_back( std::vector<Int_>() );
        _idx.back().reserve(this->_d_nz);
      }
	
#pragma omp parallel for schedule(static)
      for (Int_ i=0;i<this->tnds().size();++i) {

        Map_<Id_Nd_> &schur = this->tnds()[i]->get_data()->schur();
        Mapper_<Id_Nd_> s2nz(&schur, &_nz);
        
        // does not affect idx part but dofs are already modified
        s2nz.relate();
        
        for (Int_ j=0;j<schur.get_size();++j)
          for (Int_ k=0;k<schur.id(j).get_dof();++k)
            _idx.at(i).push_back(_nz.idx(s2nz.idx(j))+k);

        Uhm_ &uhm = this->tnds()[i]->get_data()->uhm();

        // change column major to row major
        uhm.ABR().transpose();
#pragma omp critical
        {
          ierr = MatSetValues(_A, 
                              uhm.ABR().get_m(), &_idx[i][0],
                              uhm.ABR().get_n(), &_idx[i][0],
                              (const PetscScalar*)(uhm.ABR().get_flat_buffer()),
                              ADD_VALUES);CHKERRQ(ierr);
        }
        uhm.ABR().free_flat_buffer();
      }
      time_out(UHM_TIME_LEVEL_INTERNAL);

      time_in("uhm:: petsc:: assemble:: PETSC Assembly (begin - end)",UHM_TIME_LEVEL_INTERNAL);
      ierr = MatAssemblyBegin(_A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
      ierr = MatAssemblyEnd  (_A, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
      
      time_out(UHM_TIME_LEVEL_INTERNAL);
      time_out(UHM_TIME_LEVEL_FRONT);

      if (get_verbose()) {
        ierr = MatView(_A,PETSC_VIEWER_DRAW_WORLD);CHKERRQ(ierr);
      }

      return true;
    }
  };
}


#endif
